#include "utils.h"
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <stdio.h>
#include <vector>

// Block 配置
#define BLOCK_SIZE_1D 256
#define BLOCK_SIZE_2D 16

// Kernel 1: 创建掩码（非纯白为前景）
__global__ void kernel_create_mask(const uchar4* srcImg, unsigned char* mask, int numPixels) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= numPixels) return;
    uchar4 pixel = srcImg[idx];
    mask[idx] = (pixel.x < 255 || pixel.y < 255 || pixel.z < 255) ? 1 : 0;
}

// Kernel 2: 分类内部像素和边界
__global__ void kernel_classify_pixels(
    const unsigned char* mask,
    unsigned char* strictInterior,
    unsigned char* border,
    int rows, int cols
) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    int r = blockIdx.y * blockDim.y + threadIdx.y;
    if (r == 0 || r >= rows-1 || c == 0 || c >= cols-1) return;
    int idx = r * cols + c;
    if (!mask[idx]) {
        strictInterior[idx] = 0;
        border[idx] = 0;
        return;
    }
    bool up    = mask[(r-1) * cols + c];
    bool down  = mask[(r+1) * cols + c];
    bool left  = mask[r * cols + c - 1];
    bool right = mask[r * cols + c + 1];
    if (up && down && left && right) {
        strictInterior[idx] = 1;
        border[idx] = 0;
    } else {
        strictInterior[idx] = 0;
        border[idx] = 1;
    }
}

// Kernel 3: 分离通道 + 初始化 f = source
__global__ void kernel_split_and_init(
    const uchar4* src, const uchar4* dst,
    float* r_src, float* g_src, float* b_src,
    float* r_dst, float* g_dst, float* b_dst,
    float* f_r, float* f_g, float* f_b,
    int numPixels
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= numPixels) return;
    uchar4 s = src[idx], d = dst[idx];
    r_src[idx] = static_cast<float>(s.x);
    g_src[idx] = static_cast<float>(s.y);
    b_src[idx] = static_cast<float>(s.z);
    r_dst[idx] = static_cast<float>(d.x);
    g_dst[idx] = static_cast<float>(d.y);
    b_dst[idx] = static_cast<float>(d.z);
    f_r[idx] = r_src[idx];
    f_g[idx] = g_src[idx];
    f_b[idx] = b_src[idx];
}

// Kernel 4: Jacobi 迭代（RGB 合并）
__global__ void kernel_jacobi_iteration_interior_rgb(
    float* dstR, float* dstG, float* dstB,
    const unsigned char* strictInterior,
    const float* srcR, const float* srcG, const float* srcB,
    const float* f_r, const float* f_g, const float* f_b,
    float* f_r_next, float* f_g_next, float* f_b_next,
    int rows, int cols,
    const int* x_coords, const int* y_coords,
    int num_interior
){
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_interior) return;
    int c = x_coords[idx];
    int r = y_coords[idx];
    int img_idx = r * cols + c;

    float sum_r = 0.0f, sum_g = 0.0f, sum_b = 0.0f;
    float sum_r_diff = 0.0f, sum_g_diff = 0.0f, sum_b_diff = 0.0f;

    int neighbors[4] = {img_idx - 1, img_idx + 1, img_idx - cols, img_idx + cols};
    for (int i = 0; i < 4; ++i) {
        int n_idx = neighbors[i];
        // 只判断是否在图像范围内和是否为严格内部
        bool is_valid_interior = (n_idx >= 0 && n_idx < rows * cols) && strictInterior[n_idx];
        
        // Sum1: 内部邻居用 f，边界/外部邻居用 dst
        sum_r += is_valid_interior ? f_r[n_idx] : dstR[n_idx];
        sum_g += is_valid_interior ? f_g[n_idx] : dstG[n_idx];
        sum_b += is_valid_interior ? f_b[n_idx] : dstB[n_idx];
        
        // Sum2: 源图像梯度差
        sum_r_diff += (srcR[img_idx] - srcR[n_idx]);
        sum_g_diff += (srcG[img_idx] - srcG[n_idx]);
        sum_b_diff += (srcB[img_idx] - srcB[n_idx]);
    }

    float new_r = (sum_r + sum_r_diff) / 4.0f;
    float new_g = (sum_g + sum_g_diff) / 4.0f;
    float new_b = (sum_b + sum_b_diff) / 4.0f;

    f_r_next[img_idx] = fminf(255.0f, fmaxf(0.0f, new_r));
    f_g_next[img_idx] = fminf(255.0f, fmaxf(0.0f, new_g));
    f_b_next[img_idx] = fminf(255.0f, fmaxf(0.0f, new_b));
}

// Kernel 5: 写回结果
__global__ void kernel_write_result(
    const float* red_result,
    const float* green_result,
    const float* blue_result,
    uchar4* output,
    const unsigned char* strictInterior,
    int rows, int cols
) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    int r = blockIdx.y * blockDim.y + threadIdx.y;
    if (r >= rows || c >= cols) return;
    int idx = r * cols + c;
    if (strictInterior[idx]) {
        output[idx].x = static_cast<unsigned char>(red_result[idx]);
        output[idx].y = static_cast<unsigned char>(green_result[idx]);
        output[idx].z = static_cast<unsigned char>(blue_result[idx]);
    }
}

// 主函数
void poisson_blending(const uchar4* const h_sourceImg,
                    const size_t numRowsSource, const size_t numColsSource,
                    const uchar4* const h_destImg,
                    uchar4* const h_blendedImg) {
    const size_t srcSize = numRowsSource * numColsSource;

    // 1. 初始化输出 = 目标图像
    checkCudaErrors(cudaMemcpy(h_blendedImg, h_destImg, srcSize * sizeof(uchar4), cudaMemcpyHostToHost));

    // 2. 分配设备内存
    uchar4 *d_sourceImg = nullptr, *d_destImg = nullptr, *d_blendedImg = nullptr;
    unsigned char *d_mask = nullptr, *d_strictInterior = nullptr, *d_border = nullptr;
    float *d_red_src = nullptr, *d_green_src = nullptr, *d_blue_src = nullptr;
    float *d_red_dst = nullptr, *d_green_dst = nullptr, *d_blue_dst = nullptr;
    float *d_f_red_1 = nullptr, *d_f_red_2 = nullptr;
    float *d_f_green_1 = nullptr, *d_f_green_2 = nullptr;
    float *d_f_blue_1 = nullptr, *d_f_blue_2 = nullptr;
    int *d_x_coords = nullptr, *d_y_coords = nullptr;

    // 分配所有内存
    checkCudaErrors(cudaMalloc(&d_sourceImg, srcSize * sizeof(uchar4)));
    checkCudaErrors(cudaMalloc(&d_destImg,   srcSize * sizeof(uchar4)));
    checkCudaErrors(cudaMalloc(&d_blendedImg, srcSize * sizeof(uchar4)));
    checkCudaErrors(cudaMalloc(&d_mask, srcSize * sizeof(unsigned char)));
    checkCudaErrors(cudaMalloc(&d_strictInterior, srcSize * sizeof(unsigned char)));
    checkCudaErrors(cudaMalloc(&d_border, srcSize * sizeof(unsigned char)));
    checkCudaErrors(cudaMalloc(&d_red_src, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_green_src, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_blue_src, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_red_dst, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_green_dst, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_blue_dst, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_f_red_1, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_f_red_2, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_f_green_1, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_f_green_2, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_f_blue_1, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_f_blue_2, srcSize * sizeof(float)));
    checkCudaErrors(cudaMalloc(&d_x_coords, srcSize * sizeof(int)));
    checkCudaErrors(cudaMalloc(&d_y_coords, srcSize * sizeof(int)));

    // 3. 拷贝输入数据
    checkCudaErrors(cudaMemcpy(d_sourceImg, h_sourceImg, srcSize * sizeof(uchar4), cudaMemcpyHostToDevice));
    checkCudaErrors(cudaMemcpy(d_destImg, h_destImg, srcSize * sizeof(uchar4), cudaMemcpyHostToDevice));
    checkCudaErrors(cudaMemcpy(d_blendedImg, h_blendedImg, srcSize * sizeof(uchar4), cudaMemcpyHostToDevice));

    // 4. 创建掩码
    dim3 block1d(BLOCK_SIZE_1D);
    dim3 grid1d((srcSize + BLOCK_SIZE_1D - 1) / BLOCK_SIZE_1D);
    kernel_create_mask<<<grid1d, block1d>>>(d_sourceImg, d_mask, srcSize);
    checkCudaErrors(cudaDeviceSynchronize());
    checkCudaErrors(cudaGetLastError());

    // 5. 分类内部/边界
    dim3 block2d(BLOCK_SIZE_2D, BLOCK_SIZE_2D);
    dim3 grid2d((numColsSource + BLOCK_SIZE_2D - 1) / BLOCK_SIZE_2D,
                (numRowsSource + BLOCK_SIZE_2D - 1) / BLOCK_SIZE_2D);
    kernel_classify_pixels<<<grid2d, block2d>>>(d_mask, d_strictInterior, d_border, numRowsSource, numColsSource);
    checkCudaErrors(cudaDeviceSynchronize());
    checkCudaErrors(cudaGetLastError());

    // 6. CPU 收集内部像素坐标
    std::vector<int> h_x_coords, h_y_coords;
    h_x_coords.reserve(srcSize / 4);
    h_y_coords.reserve(srcSize / 4);

    unsigned char* h_strictInterior = new unsigned char[srcSize];
    checkCudaErrors(cudaMemcpy(h_strictInterior, d_strictInterior, srcSize * sizeof(unsigned char), cudaMemcpyDeviceToHost));

    for (int r = 1; r < numRowsSource - 1; ++r) {
        for (int c = 1; c < numColsSource - 1; ++c) {
            int idx = r * numColsSource + c;
            if (h_strictInterior[idx]) {
                h_x_coords.push_back(c);
                h_y_coords.push_back(r);
            }
        }
    }
    delete[] h_strictInterior;

    int num_interior = h_x_coords.size();

    // 7. 仅在有内部像素时进行迭代
    if (num_interior > 0) {
        // 拷贝坐标
        checkCudaErrors(cudaMemcpy(d_x_coords, h_x_coords.data(), num_interior * sizeof(int), cudaMemcpyHostToDevice));
        checkCudaErrors(cudaMemcpy(d_y_coords, h_y_coords.data(), num_interior * sizeof(int), cudaMemcpyHostToDevice));

        // 分离通道 + 初始化
        kernel_split_and_init<<<grid1d, block1d>>>(
            d_sourceImg, d_destImg,
            d_red_src, d_green_src, d_blue_src,
            d_red_dst, d_green_dst, d_blue_dst,
            d_f_red_1, d_f_green_1, d_f_blue_1,
            srcSize);
        checkCudaErrors(cudaDeviceSynchronize());
        checkCudaErrors(cudaGetLastError());

        // Jacobi 迭代
        const size_t numIterations = 800;
        float *d_f_red_A = d_f_red_1,   *d_f_red_B = d_f_red_2;
        float *d_f_green_A = d_f_green_1, *d_f_green_B = d_f_green_2;
        float *d_f_blue_A = d_f_blue_1,   *d_f_blue_B = d_f_blue_2;

        dim3 grid_interior((num_interior + BLOCK_SIZE_1D - 1) / BLOCK_SIZE_1D);

        for (size_t iter = 0; iter < numIterations; ++iter) {
            if (iter % 2 == 0) {
                kernel_jacobi_iteration_interior_rgb<<<grid_interior, block1d>>>(
                    d_red_dst, d_green_dst, d_blue_dst,
                    d_strictInterior,
                    d_red_src, d_green_src, d_blue_src,
                    d_f_red_A, d_f_green_A, d_f_blue_A,
                    d_f_red_B, d_f_green_B, d_f_blue_B,
                    numRowsSource, numColsSource, d_x_coords, d_y_coords, num_interior);
            } else {
                kernel_jacobi_iteration_interior_rgb<<<grid_interior, block1d>>>(
                    d_red_dst, d_green_dst, d_blue_dst,
                    d_strictInterior,
                    d_red_src, d_green_src, d_blue_src,
                    d_f_red_B, d_f_green_B, d_f_blue_B,
                    d_f_red_A, d_f_green_A, d_f_blue_A,
                    numRowsSource, numColsSource, d_x_coords, d_y_coords, num_interior);
            }
        }
        checkCudaErrors(cudaDeviceSynchronize());
        checkCudaErrors(cudaGetLastError());

        // 确定最终结果缓冲区
        float *d_final_red   = (numIterations % 2 == 0) ? d_f_red_A : d_f_red_B;
        float *d_final_green = (numIterations % 2 == 0) ? d_f_green_A : d_f_green_B;
        float *d_final_blue  = (numIterations % 2 == 0) ? d_f_blue_A : d_f_blue_B;

        // 写回结果
        kernel_write_result<<<grid2d, block2d>>>(
            d_final_red, d_final_green, d_final_blue,
            d_blendedImg,
            d_strictInterior, numRowsSource, numColsSource);
        checkCudaErrors(cudaDeviceSynchronize());
        checkCudaErrors(cudaGetLastError());

        // 拷贝结果回主机
        checkCudaErrors(cudaMemcpy(h_blendedImg, d_blendedImg, srcSize * sizeof(uchar4), cudaMemcpyDeviceToHost));
    }

    // 8. 释放所有设备内存
    #define safeFree(p) if(p) { cudaFree(p); p = nullptr; }
    safeFree(d_sourceImg);
    safeFree(d_destImg);
    safeFree(d_blendedImg);
    safeFree(d_mask);
    safeFree(d_strictInterior);
    safeFree(d_border);
    safeFree(d_red_src);
    safeFree(d_green_src);
    safeFree(d_blue_src);
    safeFree(d_red_dst);
    safeFree(d_green_dst);
    safeFree(d_blue_dst);
    safeFree(d_f_red_1);
    safeFree(d_f_red_2);
    safeFree(d_f_green_1);
    safeFree(d_f_green_2);
    safeFree(d_f_blue_1);
    safeFree(d_f_blue_2);
    safeFree(d_x_coords);
    safeFree(d_y_coords);
}

// 接口函数
void your_blend(const uchar4* const h_sourceImg,
                const size_t numRowsSource, const size_t numColsSource,
                const uchar4* const h_destImg,
                uchar4* const h_blendedImg) {
    poisson_blending(h_sourceImg, numRowsSource, numColsSource, h_destImg, h_blendedImg);
}